/**
 * 
 */
package tls;
import java.security.InvalidKeyException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.Key; 

import javax.crypto.KeyGenerator;
import javax.crypto.Mac;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

import test.Convert;
import test.exception.*;

/**
 * @author lcym
 *
 */
public class PRFUtils {
	private byte [] last_hmacMD5;
	private byte [] last_hmacSHA1;
	private static int PRF_RESULT_LEN = 80; //bytes
	private static int HMAC_MD5_LEN = 16; //bytes
	private static int HMAC_SHA1_LEN = 20; // bytes

	/**
	 * @param args
	 */
	public static void main(String[] args) {
		// TODO Auto-generated method stub

	}
	
	/**
	 * 
	 * @param key I guess this is the PSK key
	 * @param seed
	 * @return
	 */
	public byte[] TLS_PRF(byte[] key, byte [] seed){
		String HmacMD5_tmp = "";
		String HmacSHA1_tmp = "";
		byte [] md5_key = new byte[key.length / 2];
		byte [] sha1_key = new byte[key.length / 2];
		byte [] result;
		byte [] hmac_md5 = null;
		byte [] hmac_sha1 = null;
		
		System.arraycopy(key, 0, md5_key, 0, key.length / 2);
		System.arraycopy(key, key.length / 2,sha1_key,0,key.length / 2);
		
		for(int md5_cnt = 0; md5_cnt < PRF_RESULT_LEN / HMAC_MD5_LEN; 
				md5_cnt ++){
			HmacMD5_tmp += Convert.byteArrayToHexStr(P_MD5(md5_key,seed));
		}
		
		for(int sha1_cnt = 0; sha1_cnt < PRF_RESULT_LEN / HMAC_SHA1_LEN; 
				sha1_cnt ++){
			HmacSHA1_tmp += Convert.byteArrayToHexStr(P_SHA1(sha1_key,seed));
		}
		
		try{
			hmac_md5 = Convert.hexStrToByteArray(HmacMD5_tmp);
			hmac_sha1 = Convert.hexStrToByteArray(HmacSHA1_tmp);
		}catch(EncodeDecodeException ex){
			
		}		
		result = ArrayUtils.xorArr(hmac_md5, hmac_sha1);
		
		return result;
	}
	
	/**
	 * 
	 * @param key
	 * @param seed
	 * @return
	 */
	public byte[] P_MD5(byte [] key, byte [] seed){
		byte [] cur_seed = ArrayUtils.concat(last_hmacMD5, seed);
		byte [] ret = getHmac(key,cur_seed,"HmacMD5");
		last_hmacMD5 = ret;
		return ret;
	}
	
	/**
	 * 
	 * @param key
	 * @param seed
	 * @return
	 */
	public byte[] P_SHA1(byte [] key, byte [] seed){
		byte [] cur_seed = ArrayUtils.concat(last_hmacSHA1, seed);
		byte [] ret = getHmac(key,cur_seed,"HmacSHA1");
		last_hmacSHA1 = ret;
		return ret;
	}
	
	/**
	 * 
	 * @param input
	 * @return
	 */
	public static byte [] getMD5(byte[] input){
		MessageDigest m;
		try{
			m=MessageDigest.getInstance("MD5");
			m.update(input,0,input.length);
		}catch(NoSuchAlgorithmException ex){
			
			return null;
		}
		
		return m.digest();
	}
	
	public static byte[] getHmac(byte [] privateKey, byte[] input, String sAlgo) {
		byte[] digest = null;
		
		try {
		    // Generate a key for the HMAC-MD5 keyed-hashing algorithm; see RFC 2104
		    // In practice, you would save this key.
			//TODO: consider removing the keyGen
		    KeyGenerator keyGen = KeyGenerator.getInstance(sAlgo);
		    SecretKeySpec secret = new SecretKeySpec(privateKey,sAlgo);

		    // Create a MAC object using HMAC-MD5 and initialize with key
		    Mac mac = Mac.getInstance(secret.getAlgorithm());
		    try {
				mac.init(secret);
			} catch (InvalidKeyException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}

		    // Encode the string into bytes using utf-8 and digest it
		    digest = mac.doFinal(input);

		} catch (NoSuchAlgorithmException e) {
			
		} 
		
		return digest;
	}

}
